﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Diagnostics;
using System.Collections;

namespace AndersLiu.PE
{
    public interface IImageData
    {
        PEFile PEFile { get; }
        uint FileOffset { get; }
        uint RawDataSize { get; }

        byte[] GetRawData();
    }

    public abstract class ImageData : IImageData
    {
        public PEFile PEFile { get; private set; }
        public uint FileOffset { get; private set; }
        public uint RawDataSize { get; private set; }

        public byte[] GetRawData()
        {
            return this.PEFile.GetRawData(this);
        }

        internal ImageData(PEFile image, uint offset)
        {
            Debug.Assert(image != null);
            Debug.Assert(offset < image.InternalImageBuffer.Length);

            this.PEFile = image;
            this.FileOffset = offset;
        }

        internal ImageData(PEFile image, uint offset, uint size)
            : this(image, offset)
        {
            this.RawDataSize = size;
        }

        protected internal void SetSize(uint size)
        {
            this.RawDataSize = size;
        }
    }

    public abstract class Field<T> : ImageData
    {
        public T Value
        {
            get
            {
                if (!_valueReady)
                {
                    _value = ParseValue(GetRawData());
                    _valueReady = true;
                }

                return _value;
            }
        }

        public override string ToString()
        {
            return string.Format("[{0:X8}h:{1:X}h({1})] {2}",
                this.FileOffset, this.RawDataSize, this.GetValueString());
        }

        internal Field(PEFile image, uint offset, uint size)
            : base(image, offset, size)
        {
        }

        protected abstract T ParseValue(byte[] rawData);

        protected virtual string GetValueString()
        {
            return this.Value.ToString();
        }

        private bool _valueReady = false;
        private T _value;
    }

    public class U1Field : Field<byte>
    {
        internal U1Field(PEFile image, uint offset)
            : base(image, offset, 1)
        {
        }

        protected override byte ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == 1);

            return rawData[0];
        }

        protected override string GetValueString()
        {
            return string.Format("{0:X2}h({0})", this.Value);
        }
    }

    public class U2Field : Field<ushort>
    {
        internal U2Field(PEFile image, uint offset)
            : base(image, offset, 2)
        {
        }

        protected override ushort ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == 2);

            return BitConverter.ToUInt16(rawData, 0);
        }

        protected override string GetValueString()
        {
            return string.Format("{0:X4}h({0})", this.Value);
        }
    }

    public class U4Field : Field<uint>
    {
        internal U4Field(PEFile image, uint offset)
            : base(image, offset, 4)
        {
        }

        protected override uint ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == 4);

            return BitConverter.ToUInt32(rawData, 0);
        }

        protected override string GetValueString()
        {
            return string.Format("{0:X8}h({0})", this.Value);
        }
    }

    public class U8Field : Field<ulong>
    {
        internal U8Field(PEFile image, uint offset)
            : base(image, offset, 8)
        {
        }

        protected override ulong ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == 8);

            return BitConverter.ToUInt64(rawData, 0);
        }

        protected override string GetValueString()
        {
            return string.Format("{0:X16}h({0})", this.Value);
        }
    }

    public class FixedSizeByteArrayField : Field<byte[]>
    {
        internal FixedSizeByteArrayField(PEFile image, uint offset, uint size)
            : base(image, offset, size)
        {
        }

        protected override byte[] ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == this.RawDataSize);

            return rawData;
        }

        protected override string GetValueString()
        {
            return BitConverter.ToString(this.Value);
        }
    }

    public class FixedSizeAsciiStringField : Field<string>
    {
        internal FixedSizeAsciiStringField(PEFile image, uint offset, uint size)
            : base(image, offset, size)
        {
        }

        protected override string ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == this.RawDataSize);

            var len = 0;
            while (this.PEFile.InternalImageBuffer[this.FileOffset + len] != 0)
                len++;

            return Encoding.ASCII.GetString(rawData, 0, len);
        }

        protected override string GetValueString()
        {
            return this.Value;
        }
    }

    public class FixedSizeUtf8StringField : Field<string>
    {
        internal FixedSizeUtf8StringField(PEFile image, uint offset, uint size)
            : base(image, offset, size)
        {
        }

        protected override string ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == this.RawDataSize);

            var len = 0;
            while (this.PEFile.InternalImageBuffer[this.FileOffset + len] != 0)
                len++;

            return Encoding.UTF8.GetString(rawData, 0, len);
        }

        protected override string GetValueString()
        {
            return this.Value;
        }
    }

    public class FixedSizeUnicodeStringField : Field<string>
    {
        internal FixedSizeUnicodeStringField(PEFile image, uint offset, uint size)
            : base(image, offset, size)
        {
        }

        protected override string ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == this.RawDataSize);

            return Encoding.Unicode.GetString(rawData);
        }

        protected override string GetValueString()
        {
            return this.Value;
        }
    }

    public class NullTerminatedUtf8StringField : Field<string>
    {
        internal NullTerminatedUtf8StringField(PEFile image, uint offset)
            : base(image, offset, 0)
        {
            var len = (uint)0;
            while (this.PEFile.InternalImageBuffer[this.FileOffset + len] != 0)
                len++;

            SetSize(len + 1);  // includes the null terminator
        }

        protected override string ParseValue(byte[] rawData)
        {
            Debug.Assert(rawData != null && rawData.Length == this.RawDataSize);

            return Encoding.UTF8.GetString(rawData, 0, (int)(this.RawDataSize - 1));
        }

        protected override string GetValueString()
        {
            return this.Value;
        }
    }

    public class RvaField : U4Field
    {
        public bool IsEmpty
        {
            get
            {
                return this.Value == 0;
            }
        }

        public uint Offset
        {
            get
            {
                return this.PEFile.RvaToOffset(this.Value);
            }
        }

        public uint SectionOffset
        {
            get
            {
                return this.Value - this.Section.VirtualAddress.Value;
            }
        }

        public ImageSectionHeader Section
        {
            get
            {
                return this.PEFile.GetSection(this.Value);
            }
        }

        internal RvaField(PEFile image, uint offset)
            : base(image, offset)
        {
        }

        protected override string GetValueString()
        {
            if (this.IsEmpty)
            {
                return "[RVA:0]";
            }
            else
            {
                return string.Format("[RVA:{0:X8}h, Offset:{1:X8}h, {2}+{3:X}h]",
                    this.Value, this.Offset, this.Section.Name.Value, this.SectionOffset);
            }
        }
    }

    public class GuidField : Field<Guid>
    {
        internal GuidField(PEFile image, uint offset)
            : base(image, offset, 16)
        {
        }

        protected override Guid ParseValue(byte[] rawData)
        {
            return new Guid(rawData);
        }

        protected override string GetValueString()
        {
            return this.Value.ToString();
        }
    }

    public class CompressedUIntField : Field<uint>
    {
        internal CompressedUIntField(PEFile image, uint offset)
            : base(image, offset, IntUtility.GetCompressedIntSize(image.InternalImageBuffer[offset]))
        {
        }

        protected override uint ParseValue(byte[] rawData)
        {
            return IntUtility.DecompressUInt(rawData);
        }

        protected override string GetValueString()
        {
            return string.Format("{0:X8}h({0}) <{1}>", this.Value, BitConverter.ToString(GetRawData()));
        }
    }

    public class U4Vector : ImageData, IEnumerable<U4Field>
    {
        public U4Field this[int index] { get { return _u4s[index]; } }

        public int Count { get { return _u4s.Length; } }

        public IEnumerator<U4Field> GetEnumerator()
        {
            for (var i = 0; i < _u4s.Length; i++)
                yield return _u4s[i];
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return this.GetEnumerator();
        }

        internal U4Vector(PEFile image, uint offset, int length)
            : base(image, offset)
        {
            Debug.Assert(length >= 0);

            var pos = offset;

            _u4s = new U4Field[length];

            for (var i = 0; i < length; i++)
            {
                _u4s[i] = new U4Field(image, pos);
                pos += _u4s[i].RawDataSize;
            }

            SetSize(pos - offset);
        }

        private U4Field[] _u4s;
    }
}
